"""\
Data preparation routine.

TODO: 
* implement multiple-exposure stitching (from matlab code)
* implement image orientation.
* improve and test dark field subtraction
* implement masks
* add spec interface
* clean up the default parameters (maybe using a beamline-dependent preprocessing routine?)  
* test
"""

import pyE17
import numpy as np
import glob
import os
import re

from pyE17 import verbose
from pyE17 import io
from pyE17 import utils as U

__all__ = ['prepare_data']

def wait_for_figure(fig):
    while fig.canvas.get_visible():
        fig.waitforbuttonpress()

default_parameters = dict(
    verbose_level = 2,
    image_orientation = [1, 1, 0], # id22 : flipped upside-down compared to Pilatus 2M at cSAXS
    user = None,
    prefix = None,
    prefix_pattern = '{0.user}_1_',
    save_prefix = None,
    save_prefix_pattern = 'S{0.scan_number:05d}',
    data_ctr = None,
    multexp = False,
    dark_frame = None,
    filename_pattern = '{0.read_path}/{0.prefix}{0.scan_number:05d}_{index}.{0.file_extension}',
    filename_multexp_pattern = '{0.read_path}/{0.prefix}{0.scan_number:05d}_{index}_{exp}.{0.file_extension}',
    write_path = None,
    read_path = None,
    write_path_pattern = '{base_path}/analysis/S{scan_number:05d}',
    read_path_pattern = '{base_path}/pilatus/S{smin:05d}-{smax:05d}/S{scan_number:05d}',
    write_path_from_scan = lambda pp, scan_number : pp.write_path_pattern.format(pp, base_path=pp.base_path, scan_number=scan_number),
    read_path_from_scan = lambda pp, scan_number: pp.read_path_pattern.format(pp, base_path=pp.base_path,smin=int(1000*np.floor(scan_number/1000.)),smax=int(1000*np.floor(scan_number/1000.)+999), scan_number=scan_number),
    save_filename_pattern = '{0.write_path}/{0.save_prefix}_data_{0.asize[0]:03d}x{0.asize[1]:03d}.h5',
    force_overwrite = False,
    fmask = None,
    count_cutoff = None
)

def prepare_data(pdict=None, **kwargs):
    """\
        Data preparation routine.
    """
    p = U.Param(kwargs,pdict,default_parameters)
    verbose.set_level(p.verbose_level)

    scan_number = p.get('scan_number', None)
    p.scan_number = scan_number
    if scan_number is not None:
        verbose(2, 'Processing scan number %d' % scan_number)
    else:
        scan_number = 1 # Fake number
        
    # Check for write path
    if p.write_path is None:
        #p.write_path = p.write_path_from_scan(p.write_path_pattern, scan_number)
        p.write_path = p.write_path_from_scan(p, scan_number)

    verbose(2, 'Will write to path: %s' % p.write_path)

    # Check for read path
    if p.read_path is None:
        #p.read_path = p.read_path_from_scan(p.read_path_pattern, scan_number)
        p.read_path = p.read_path_from_scan(p, scan_number)
  
    verbose(2, 'Will read from path: %s' % p.read_path)

    # Prefix
    if p.prefix is None:
        p.prefix = p.prefix_pattern.format(p)

    if p.save_prefix is None:
        p.save_prefix = p.save_prefix_pattern.format(p)


    # It seems safe to always override the file extension since image_read will take care of that.
    p.file_extension = '*'

    if p.multexp:
        raise RuntimeError('Multiple exposure scans are not yet supported')

    verbose(3, 'Looking for a first file...')
    
    # Play some tricks on the pattern string to get the proper wildcard string to glob
    # pattern_spec = dict([(k,int(v)) for k,v in re.findall('\{([^:]+):([0-9]*)[a-z]\}', p.filename_pattern)])
    # filename_pattern_no_spec = re.sub('\{([^:]+)[^\}]*\}', r'{\1}', p.filename_pattern)
    try:
        index_length = int(re.findall('\{index:([0-9]*)[a-z]\}', p.filename_pattern)[0])
    except ValueError:
        print 'filename_pattern should have a substring "{index:?d}"'
        raise
    filename_pattern_no_spec = re.sub('\{index[^\}]*\}', '{index}', p.filename_pattern)
    fpattern_wildcard = filename_pattern_no_spec.format(p, index='?'*index_length)
    file_list = glob.glob(fpattern_wildcard)
    if file_list:
        verbose(3, 'Found %d files matching "%s"' % (len(file_list), fpattern_wildcard))
    else:
        raise IOError('No file matching "%s"' % fpattern_wildcard)

    if p.count_cutoff is not None:
        count_cutoff = p.count_cutoff
        verbose(2, 'Intensty cutoff : %f' % count_cutoff)
    else:
        count_cutoff = np.inf

    multexp_stuff = """\
        # multiple exposures
        fpattern = p.filename_multexp_pattern.format(p, exp='*', index='{index}')
        this_fp = fpattern.format(index='00000')
        file_list = glob.glob(this_fp)
        p.multexp = True 
        verbose(3, 'Found a file matching "%s"' % this_fp)
        if not file_list:
            raise IOError('No file match!')
        else:
            verbose(2, 'This is a multiple-exposure dataset')
    """

    # Read in a first dataset to evaluate the size and (eventually) the center.
    first_filename = sorted(file_list)[0]
    f,meta = io.image_read(first_filename, doglob=True)

    multexp_stuff = """\
    if num_exp > 1:
    
        exp_times = np.array([float(m['time_of_frame']) for m in meta])

        min_exp = exp_times.min()
        low_exp_index = exp_times.argmin()

        all_same = False
        if min_exp == exp_times.mean():
            # all exposure times are the same
            all_same = True
            verbose(2, 'Exposure time: %f' % min_exp)
        else:
            max_exp = exp_times.max()
            high_exp_index = exp_times.argmax()
            verbose(2, 'Lowest exposure time: %f (number %d)' % (min_exp, low_exp_index))
            verbose(2, 'Highest exposure time: %f (number %d)' % (max_exp, high_exp_index))
    """
    f = f[0]
    sh = f.shape

    if p.fmask is None:
        fmask = np.ones_like(f)
        verbose(3, 'No fmask - using ones.')
    else:
        verbose(3, 'Using provided fmask')
        fmask = p.fmask
    # pilatus-specific - should be dealt with in a more elegant way
    """\
    % if all(sz == [407 487])
    %     mask = mask(637:1043, 495:981);
    % 
    % %--- EDIT: CMK 11.12.2008 ---
    % % Included option for 4-bank readout.    
    % elseif all(sz == [831 1475])
    %     mask = mask(425:1255, 1:1475);
    %     size(mask)
    % %--- end of edit ---
    % 
    % elseif ~all(sz == [1679 1475])
    %     error('Unsupported frame dimensions : %d x %d', sz(1), sz(2));
    % end
    """

    # Check array size
    if not p.hasattr('asize'):
        raise RuntimeError('Array size (asize) has not been defined')

    asize = p.asize
 
    fullframe = False
    if asize is None:
        asize = sh
        verbose(2, 'Full frames (%d x %d) will be saved (so no recentering).' % (f.shape))
        fullframe = True        
    elif np.isscalar(asize):
            asize = (asize,asize)
    p.asize = asize
    asize = np.array(asize)

    # Save filename
    data_filename = p.save_filename_pattern.format(p)
    verbose(2,'Saving dataset to file %s ' % data_filename)
    if not os.path.exists(p.write_path):
        os.makedirs(p.write_path)
    if os.path.exists(data_filename):
        if p.force_overwrite == True:
            verbose(1, 'Save file exists but will be overwritten (force_overwrite is True)')
        elif p.force_overwrite == False:
            raise RuntimeError('File %s exists! Operation cancelled.' % data_filename)
        elif p.force_overwrite == None:
            ans = raw_input('File %s exists! Overwrite? [Y]/N' % data_filename)
            if ans and ans.upper() != 'Y':
                raise RuntimeError('Operation cancelled by user.')


    if not fullframe:
        # Compute center of mass
        cx = (f*fmask).sum(axis=0).argmax()
        cy = (f*fmask).sum(axis=1).argmax()
        ctr_auto = (cy, cx)

        # Check for center position
        ctr = p.data_ctr
        if ctr is None:
            ctr = ctr_auto
            verbose(2, 'Using center: (%d, %d)' % ctr)
        elif ctr == 'inter':
            import matplotlib as mpl
            fig = mpl.pyplot.figure()
            ax = fig.add_subplot(1,1,1)
            ax.imshow(np.log(f))
            ax.set_title('Select center point (hit return to finish)')
            s = U.Multiclicks(ax, True, mode='replace')
            mpl.pyplot.show()
            s.wait_until_closed()
            ctr = np.round(np.array(s.pts[0][::-1]));
            verbose(2, 'Using center: (%d, %d) - I would have guessed it is (%d, %d)' % (ctr[0], ctr[1], ctr_auto[0], ctr_auto[1]))
        else:
            verbose(2, 'Using center: (%d, %d) - I would have guessed it is (%d, %d)' % (ctr[0], ctr[1], ctr_auto[0], ctr_auto[1]))

        ctr = np.array(ctr)
        lim_inf = ctr - asize/2.
        lim_sup = ctr + asize/2.
        if (lim_inf < 0).any() or (lim_sup >= np.array(sh)).any():
            verbose(1, 'Warning: chosen center is too close to the edge! Changing the coordinates to make it fit.')
            out_string = 'From ' + str(ctr)
            ctr -= lim_inf * (lim_inf < 0)
            ctr -= (lim_sup - np.array(sh)) * (lim_sup >= np.array(sh))
            lim_inf = ctr - asize/2.
            lim_sup = ctr + asize/2.
            out_string += ' to ' + str(ctr)
            verbose(1, out_string)

        p.data_ctr = ctr
        fmask = fmask[lim_inf[0]:lim_sup[0], lim_inf[1]:lim_sup[1]]

    # Prepare the general meta-information dictionnary
    meta_list = dict([(k, []) for k in meta[0].keys()])

    numpts = np.prod(asize)

    if fullframe:
        f,meta = io.image_read(fpattern_wildcard, doglob=True)
    else:
        f,meta = io.image_read(fpattern_wildcard, doglob=True, roi=(lim_inf[0],lim_sup[0],lim_inf[1],lim_sup[1]))

    for mm in meta:
        for k,v in mm.items():
            meta_list[k].append(v)

    npts = len(f)
    verbose(2, 'Read %d files' % npts)
    data = np.zeros((npts, asize[0], asize[1]), dtype=np.single);
    if p.dark_frame is not None:
        for nn in range(npts):
            newframe = f[nn] - p.dark_frame
            newframe[newframe < 0] = 0
            data[nn,:,:] = newframe
    else:
        for nn in range(npts):
            data[nn,:,:] = f[nn]

    # This makes sure that anything in the data structure that cannot be saved will just be ignored.
    h5opt = pyE17.io.h5options['UNSUPPORTED']
    pyE17.io.h5options['UNSUPPORTED'] = 'ignore'

    # Save data, mask, AND parameter dictionary
    pyE17.io.h5write(data_filename, data=data, fmask=fmask, paramdict=p.paramdict, meta=meta_list);

    # Restore previous behavior of h5write
    pyE17.io.h5options['UNSUPPORTED'] = h5opt

    return data_filename
